import os
import json
import glob
import loguru
from   typing   import List, Union
from   datasets import load_dataset, load_from_disk

data_path_prefix = "/home/yangxianpku/datasets/"
data_dict        = ["ShengbinYue/DISC-Law-SFT", ]

logger = loguru.logger

def datasets_download():
    try:
        for data_file in data_dict:
            dataset = load_dataset(data_file)
            dataset.save_to_disk(os.path.join(data_path_prefix, data_file))
        return True
    except Exception as e:
        logger.exception(f"Error downloading: {e}")
        return False


def datasets_format_convert(dataset, files:Union[str, List[str]]=None, output_file:str=None):
    json_data = []
    if files is None:
        file_list = glob.glob(os.path.join(data_path_prefix, dataset, '*.jsonl'))
    else:
        if isinstance(files, str):
            files = [files]
        file_list = [os.path.join(data_path_prefix, dataset, file) for file in files]

    for file in file_list:
        with open(file, 'r', encoding='utf-8') as f:
            for line in f.readlines():
                json_data.append(json.loads(line))

    template = []
    for idx, data in enumerate(json_data):
        format_json = [
                {
                    "from"  : "user",
                    "value" : data["input"]
                },
                {
                    "from"  : "assistant",
                    "value" : data["output"]
                }
            ]
        template.append({
            "id" : f"identity_{idx}",
            "conversations" : format_json
        })
    logger.info(f"微调数据总数: {len(template)}")
    
    output_file = "sft_data.json" if output_file is None else output_file
    with open(os.path.join(output_file), 'w', encoding='utf-8') as f:
        json.dump(template, f, ensure_ascii=False, indent=2)
    logger.info(f"格式化后的数据已保存至: {output_file}")
    

if __name__ == "__main__":
    # datasets_download()
    datasets_format_convert(dataset='DISC-Law-SFT',
            files="DISC-Law-SFT-Triplet-released.jsonl")